import numpy as np

def generate_anchors(base_size=16, ratios=None, scales=None):
    """
    Generate anchor (reference) windows by enumerating aspect ratios X
    scales w.r.t. a reference window.
    """

    if ratios is None:
        ratios = np.array([0.5, 1, 2])

    if scales is None:
        scales = np.array([2 ** 0, 2 ** (1.0 / 3.0), 2 ** (2.0 / 3.0)])

    num_anchors = len(ratios) * len(scales)

    # initialize output anchors
    anchors = np.zeros((num_anchors, 4))

    # scale base_size
    anchors[:, 2:] = base_size * np.tile(scales, (2, len(ratios))).T

    # compute areas of anchors
    areas = anchors[:, 2] * anchors[:, 3]

    # correct for ratios
    anchors[:, 2] = np.sqrt(areas / np.repeat(ratios, len(scales)))
    anchors[:, 3] = anchors[:, 2] * np.repeat(ratios, len(scales))

    # transform from (x_ctr, y_ctr, w, h) -> (x1, y1, x2, y2)
    anchors[:, 0::2] -= np.tile(anchors[:, 2] * 0.5, (2, 1)).T
    anchors[:, 1::2] -= np.tile(anchors[:, 3] * 0.5, (2, 1)).T

    return anchors

def shift(shape, stride, anchors):
    shift_x = (np.arange(0, shape[1]) + 0.5) * stride
    shift_y = (np.arange(0, shape[0]) + 0.5) * stride

    shift_x, shift_y = np.meshgrid(shift_x, shift_y)

    shifts = np.vstack((
        shift_x.ravel(), shift_y.ravel(),
        shift_x.ravel(), shift_y.ravel()
    )).transpose()

    # add A anchors (1, A, 4) to
    # cell K shifts (K, 1, 4) to get
    # shift anchors (K, A, 4)
    # reshape to (K*A, 4) shifted anchors
    A = anchors.shape[0]
    K = shifts.shape[0]
    all_anchors = (anchors.reshape((1, A, 4)) + shifts.reshape((1, K, 4)).transpose((1, 0, 2)))
    all_anchors = all_anchors.reshape((K * A, 4))

    return all_anchors

# get the anchor information
def get_anchor():
    pyramid_levels = [3, 4, 5, 6, 7]
    strides = [2 ** x for x in pyramid_levels] # [8, 16. 32. 64. 128]
    sizes = [2 ** (x + 2) for x in pyramid_levels] # [32, 64, 128, 256, 512]
    ratios = np.array([0.5, 1, 2])
    scales = np.array([2 ** 0, 2 ** (1.0 / 3.0), 2 ** (2.0 / 3.0)]) # 1, 1.26, 1.587

    image_shape = np.array([608, 608])
    # [76 76] [38 38] [19 19] [10 10] [5 5] feature size
    image_shapes = [(image_shape + 2 ** x - 1) // (2 ** x) for x in pyramid_levels]

    all_anchors = np.zeros((0, 4)).astype(np.float32) # save anchor as [x1 y1 x2 y2]

    for idx, p in enumerate(pyramid_levels): # 5 layer
        anchors = generate_anchors(base_size=sizes[idx], ratios=ratios, scales=scales) # get 9 anchor for this layer
        shifted_anchors = shift(image_shapes[idx], strides[idx], anchors) # all anchor in this layer
        all_anchors = np.append(all_anchors, shifted_anchors, axis=0)

    return  all_anchors




# get bbox from anchor & regression
def regression(boxes, deltas):
    mean = np.array([0, 0, 0, 0]).astype(np.float32)
    std = np.array([0.1, 0.1, 0.2, 0.2]).astype(np.float32)

    widths = boxes[:, 2] - boxes[:, 0]
    heights = boxes[:, 3] - boxes[:, 1]
    ctr_x = boxes[:, 0] + 0.5 * widths
    ctr_y = boxes[:, 1] + 0.5 * heights

    dx = deltas[:, 0] * std[0] + mean[0]
    dy = deltas[:, 1] * std[1] + mean[1]
    dw = deltas[:, 2] * std[2] + mean[2]
    dh = deltas[:, 3] * std[3] + mean[3]

    pred_ctr_x = ctr_x + dx * widths
    pred_ctr_y = ctr_y + dy * heights
    pred_w = np.exp(dw) * widths
    pred_h = np.exp(dh) * heights

    pred_boxes_x1 = pred_ctr_x - 0.5 * pred_w
    pred_boxes_y1 = pred_ctr_y - 0.5 * pred_h
    pred_boxes_x2 = pred_ctr_x + 0.5 * pred_w
    pred_boxes_y2 = pred_ctr_y + 0.5 * pred_h

    pred_boxes = np.stack([pred_boxes_x1, pred_boxes_y1, pred_boxes_x2, pred_boxes_y2], 1)

    return pred_boxes

# keep bbox in img
def clipBoxes(boxes, img):
    num_channels, height, width = img.shape

    boxes[:, 0] = np.clip(boxes[:, 0], a_min=0, a_max=width)
    boxes[:, 1] = np.clip(boxes[:, 1], a_min=0, a_max=height)
    boxes[:, 2] = np.clip(boxes[:, 2], a_min=0, a_max=width)
    boxes[:, 3] = np.clip(boxes[:, 3], a_min=0, a_max=height)

    return boxes


def CIoU_nms(scores, classes, bboxes, nms_threshold=0.5, flag=0):
    sort_scores = -np.sort(-scores) # default up-sort
    sort_scores_indexes = np.argsort(-scores)
    sort_classes = classes[sort_scores_indexes]
    sort_bboxes = bboxes[sort_scores_indexes]

    sort_bboxes_wh = sort_bboxes[:, 2:] - sort_bboxes[:, :2]
    sort_bboxes_area = sort_bboxes_wh[:, 0] * sort_bboxes_wh[:, 1]
    detected_classes = np.unique(sort_classes)

    keep_scores, keep_classes, keep_bboxes = [], [], []
    for detected_class in detected_classes:
        scores_tmp = sort_scores[sort_classes == detected_class]
        bboxes_tmp = sort_bboxes[sort_classes == detected_class]
        bboxes_area_tmp = sort_bboxes_area[sort_classes == detected_class]
        classes_tmp = sort_classes[sort_classes == detected_class]

        keep_scores_tmp, keep_classes_tmp, keep_bboxes_tmp = [], [], []

        while len(scores_tmp) > 0:
            top1_score, top1_class, top1_bbox = scores_tmp[0:1], classes_tmp[0:1], bboxes_tmp[0:1]
            top1_areas = bboxes_area_tmp[0]

            keep_scores_tmp.append(top1_score)
            keep_classes_tmp.append(top1_class)
            keep_bboxes_tmp.append(top1_bbox)

            if len(scores_tmp) == 1:
                break

            scores_tmp = scores_tmp[1:]
            classes_tmp = classes_tmp[1:]
            bboxes_tmp = bboxes_tmp[1:]
            bboxes_area_tmp = bboxes_area_tmp[1:]

            overlap_area_top_left = np.maximum(bboxes_tmp[:, :2], top1_bbox[:, :2])
            overlap_area_bot_right = np.minimum(bboxes_tmp[:, 2:], top1_bbox[:, 2:])
            overlap_area_sizes = np.clip(overlap_area_bot_right - overlap_area_top_left, a_min=0, a_max=None)
            overlap_area = overlap_area_sizes[:, 0] * overlap_area_sizes[:, 1]

            union_area = top1_areas + bboxes_area_tmp - overlap_area
            union_area = np.clip(union_area, a_min=1e-4, a_max=None)
            ious = overlap_area / union_area # IoU

            top1_bbox_ctr = (top1_bbox[:, 2:4] + top1_bbox[:, 0:2]) / 2
            bboxes_tmp_ctr = (bboxes_tmp[:, 2:4] + bboxes_tmp[:, 0:2]) / 2
            p2 = (top1_bbox_ctr[:, 0] - bboxes_tmp_ctr[:, 0]) ** 2 + (top1_bbox_ctr[:, 1] - bboxes_tmp_ctr[:, 1]) ** 2

            enclose_area_top_left = np.minimum(top1_bbox[:, 0:2], bboxes_tmp[:, 0:2])
            enclose_area_bot_right = np.maximum(top1_bbox[:, 2:4], bboxes_tmp[:, 2:4])
            enclose_area_sizes = np.clip(enclose_area_bot_right - enclose_area_top_left, a_min=1e-4, a_max=None)
            c2 = (enclose_area_sizes[:, 0]) ** 2 + (enclose_area_sizes[:, 1]) ** 2

            top1_bbox_wh = top1_bbox[:, 2:4] - top1_bbox[:, 0:2]
            bboxes_tmp_wh = bboxes_tmp[:, 2: 4] - bboxes_tmp[:, 0: 2]
            import math
            tmp = np.arctan(bboxes_tmp_wh[:, 0] / bboxes_tmp_wh[:, 1]) - np.arctan(top1_bbox_wh[:, 0] / top1_bbox_wh[:, 1])
            v = np.power(tmp, 2) * (4 / (math.pi ** 2))
            alpha = v / (1 - ious + v)

            dious = ious - p2 / c2
            cious = ious - p2 / c2 - alpha * v

            if flag == 1:
                standard = dious
            elif flag == 2:
                standard = ious
            else:
                standard = cious

            scores_tmp = scores_tmp[standard < nms_threshold]
            classes_tmp = classes_tmp[standard < nms_threshold]
            bboxes_tmp = bboxes_tmp[standard < nms_threshold]
            bboxes_area_tmp = bboxes_area_tmp[standard < nms_threshold]

        keep_scores_tmp = np.concatenate(keep_scores_tmp, 0)
        keep_classes_tmp = np.concatenate(keep_classes_tmp, 0)
        keep_bboxes_tmp = np.concatenate(keep_bboxes_tmp, 0)

        keep_scores.append(keep_scores_tmp)
        keep_classes.append(keep_classes_tmp)
        keep_bboxes.append(keep_bboxes_tmp)

    keep_scores = np.concatenate(keep_scores, 0)
    keep_classes = np.concatenate(keep_classes, 0)
    keep_bboxes = np.concatenate(keep_bboxes, 0)
    return keep_scores, keep_classes, keep_bboxes



